import psp2d, os

def get_size(OPEN_FILE, BYTE_AMOUNT = 4):
	#Get the offset size from an opened file.
	#OPEN_FILE is the opened file that we want to read from.
	#BYTE_AMOUNT can be either two or four
	#depending on the offset we're working with,
	#it is four by default.

	BYTE_ONE = ord(OPEN_FILE.read(1))
	BYTE_TWO = ord(OPEN_FILE.read(1)) * 256
	if BYTE_AMOUNT == 4:
		BYTE_THREE = ord(OPEN_FILE.read(1)) * 256 * 256
		BYTE_FOUR = ord(OPEN_FILE.read(1)) * 256 * 256 * 256
		return BYTE_ONE + BYTE_TWO + BYTE_THREE + BYTE_FOUR
	return BYTE_ONE + BYTE_TWO
#End of get_size


def convert_size(INT_SIZE, BYTE_AMOUNT = 4):
	#Convert an integer into a little endian format string
	#so it can be written to the new eboot.
	#INT_SIZE is the integer that we want to convert to a string.
	#BYTE_AMOUNT can be either two or four
	#depending on the offset we're working with,
	#it is four by default.

	if BYTE_AMOUNT == 4:
		BYTE_FOUR = 0
		while INT_SIZE >= 256 * 256 * 256:
			BYTE_FOUR = BYTE_FOUR + 1
			INT_SIZE = INT_SIZE - 256 * 256 * 256
		BYTE_THREE = 0
		while INT_SIZE >= 256 * 256:
			BYTE_THREE = BYTE_THREE + 1
			INT_SIZE = INT_SIZE - 256 * 256
	BYTE_TWO = 0
	while INT_SIZE >= 256:
		BYTE_TWO = BYTE_TWO + 1
		INT_SIZE = INT_SIZE - 256
	BYTE_ONE = INT_SIZE
	if BYTE_AMOUNT == 4:
		return chr(BYTE_ONE) + chr(BYTE_TWO) + chr(BYTE_THREE) + chr(BYTE_FOUR)
	return chr(BYTE_ONE) + chr(BYTE_TWO)
#End of convert_size


def load_eboot_image(EBOOT_SOURCE,IMAGE_TYPE):
	#Load an image from an eboot.
	#EBOOT_SOURCE can either be an open file or a path to the eboot.
	#IMAGE_TYPE can be either "icon1", "pic0" or "pic1".

	if str(type(EBOOT_SOURCE)) == "<type 'str'>":
		EBOOT_FILE = open(EBOOT_SOURCE,"rb")
	elif str(type(EBOOT_SOURCE)) == "<type 'file'>":
		EBOOT_FILE = EBOOT_SOURCE

	EBOOT_FILE.seek(0)
	if EBOOT_FILE.read(4) != chr(0) + "PBP":
		return False

	if IMAGE_TYPE == "icon0":
		EBOOT_FILE.seek(12)
		EBOOT_ICON0_START = get_size(EBOOT_FILE,4)
		EBOOT_ICON0_SIZE = get_size(EBOOT_FILE,4) - EBOOT_ICON0_START
		if EBOOT_ICON0_SIZE == 0:
			EBOOT_FILE.close()
			return False
		EBOOT_FILE.seek(EBOOT_ICON0_START)
		EBOOT_ICON0_MAGIC = EBOOT_FILE.read(4)
		EBOOT_FILE.seek(EBOOT_ICON0_START)
		if EBOOT_ICON0_MAGIC != chr(137) + "PNG":
			EBOOT_FILE.close()
			return False
		ICON1_FILE = open("ms0:/eboot_image_icon0.png","wb")
		ICON1_FILE.write(EBOOT_FILE.read(EBOOT_ICON0_SIZE))
		ICON1_FILE.close()
		ICON1_IMAGE = psp2d.Image("ms0:/eboot_image_icon0.png")
		os.remove("ms0:/eboot_image_icon0.png")
		EBOOT_FILE.close()
		return ICON1_IMAGE

	if IMAGE_TYPE == "pic0":
		EBOOT_FILE.seek(20)
		EBOOT_PIC0_START = get_size(EBOOT_FILE,4)
		EBOOT_PIC0_SIZE = get_size(EBOOT_FILE,4) - EBOOT_PIC0_START
		if EBOOT_PIC0_SIZE == 0:
			EBOOT_FILE.close()
			return False
		EBOOT_FILE.seek(EBOOT_PIC0_START)
		EBOOT_PIC0_MAGIC = EBOOT_FILE.read(4)
		EBOOT_FILE.seek(EBOOT_PIC0_START)
		if EBOOT_PIC0_MAGIC != chr(137) + "PNG":
			EBOOT_FILE.close()
			return False
		PIC0_FILE = open("ms0:/eboot_image_pic0.png","wb")
		PIC0_FILE.write(EBOOT_FILE.read(EBOOT_PIC0_SIZE))
		PIC0_FILE.close()
		PIC0_IMAGE = psp2d.Image("ms0:/eboot_image_pic0.png")
		os.remove("ms0:/eboot_image_pic0.png")
		EBOOT_FILE.close()
		return PIC0_IMAGE

	if IMAGE_TYPE == "pic1":
		EBOOT_FILE.seek(24)
		EBOOT_PIC1_START = get_size(EBOOT_FILE,4)
		EBOOT_PIC1_SIZE = get_size(EBOOT_FILE,4) - EBOOT_PIC1_START
		if EBOOT_PIC1_SIZE == 0:
			EBOOT_FILE.close()
			return False
		EBOOT_FILE.seek(EBOOT_PIC1_START)
		EBOOT_PIC1_MAGIC = EBOOT_FILE.read(4)
		EBOOT_FILE.seek(EBOOT_PIC1_START)
		if EBOOT_PIC1_MAGIC != chr(137) + "PNG":
			EBOOT_FILE.close()
			return False
		PIC1_FILE = open("ms0:/eboot_image_pic1.png","wb")
		PIC1_FILE.write(EBOOT_FILE.read(EBOOT_PIC1_SIZE))
		PIC1_FILE.close()
		PIC1_IMAGE = psp2d.Image("ms0:/eboot_image_pic1.png")
		os.remove("ms0:/eboot_image_pic1.png")
		EBOOT_FILE.close()
		return PIC1_IMAGE
#End of load_image


def unpack_eboot(EBOOT_SOURCE, EBOOT_DEST):
	#Unpack an eboot.
	#EBOOT_SOURCE can be the path to the eboot or an opened file.
	#EBOOT_DEST is the folder we want to unpack the eboot to.
	#Returns true on success, otherwise it returns false.

	EBOOT_HEADER = ""
	if str(type(EBOOT_SOURCE)) == "<type 'str'>":
		EBOOT_FILE = open(EBOOT_SOURCE,"rb")
		EBOOT_HEADER = EBOOT_FILE.read(4)
	elif str(type(EBOOT_SOURCE)) == "<type 'file'>":
		EBOOT_FILE = EBOOT_SOURCE
		EBOOT_FILE.seek(0)
		EBOOT_HEADER = EBOOT_FILE.read(4)
	else:
		return False

	if EBOOT_HEADER == chr(0) + "PBP":
		EBOOT_FILE.seek(8)
		EBOOT_PARAM = get_size(EBOOT_FILE,4)
		EBOOT_ICON0 = get_size(EBOOT_FILE,4)
		EBOOT_ICON1 = get_size(EBOOT_FILE,4)
		EBOOT_PIC0 = get_size(EBOOT_FILE,4)
		EBOOT_PIC1 = get_size(EBOOT_FILE,4)
		EBOOT_SND0 = get_size(EBOOT_FILE,4)
		EBOOT_DATAPSP = get_size(EBOOT_FILE,4)
		EBOOT_DATAPSAR = get_size(EBOOT_FILE,4)
		EBOOT_FILE.seek(0,2)
		EBOOT_END = EBOOT_FILE.tell()

		EBOOT_EXTRACT = open(EBOOT_DEST + os.sep + "PARAM.SFO","wb")
		EBOOT_FILE.seek(EBOOT_PARAM)
		EBOOT_EXTRACT.write(EBOOT_FILE.read(EBOOT_ICON0 - EBOOT_PARAM))
		EBOOT_EXTRACT.close()
		if EBOOT_ICON1 - EBOOT_ICON0 > 0:
			EBOOT_EXTRACT = open(EBOOT_DEST + os.sep + "ICON0.PNG","wb")
			EBOOT_FILE.seek(EBOOT_ICON0)
			FILE_SIZE = EBOOT_ICON1 - EBOOT_ICON0
			while FILE_SIZE > 0:
				if FILE_SIZE >= 1024 * 256:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(1024 * 256))
					FILE_SIZE = FILE_SIZE - 1024 * 256
				else:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(FILE_SIZE))
					FILE_SIZE = 0
			EBOOT_EXTRACT.close()
		if EBOOT_PIC0 - EBOOT_ICON1 > 0:
			EBOOT_EXTRACT = open(EBOOT_DEST + os.sep + "ICON1.PMF","wb")
			EBOOT_FILE.seek(EBOOT_ICON1)
			FILE_SIZE = EBOOT_PIC0 - EBOOT_ICON1
			while FILE_SIZE > 0:
				if FILE_SIZE >= 1024 * 256:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(1024 * 256))
					FILE_SIZE = FILE_SIZE - 1024 * 256
				else:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(FILE_SIZE))
					FILE_SIZE = 0
			EBOOT_EXTRACT.close()
		if EBOOT_PIC1 - EBOOT_PIC0 > 0:
			EBOOT_EXTRACT = open(EBOOT_DEST + os.sep + "PIC0.PNG","wb")
			EBOOT_FILE.seek(EBOOT_PIC0)
			FILE_SIZE = EBOOT_PIC1 - EBOOT_PIC0
			while FILE_SIZE > 0:
				if FILE_SIZE >= 1024 * 256:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(1024 * 256))
					FILE_SIZE = FILE_SIZE - 1024 * 256
				else:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(FILE_SIZE))
					FILE_SIZE = 0
			EBOOT_EXTRACT.close()
		if EBOOT_SND0 - EBOOT_PIC1 > 0:
			EBOOT_EXTRACT = open(EBOOT_DEST + os.sep + "PIC1.PNG","wb")
			EBOOT_FILE.seek(EBOOT_PIC1)
			FILE_SIZE = EBOOT_SND0 - EBOOT_PIC1
			while FILE_SIZE > 0:
				if FILE_SIZE >= 1024 * 256:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(1024 * 256))
					FILE_SIZE = FILE_SIZE - 1024 * 256
				else:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(FILE_SIZE))
					FILE_SIZE = 0
			EBOOT_EXTRACT.close()
		if EBOOT_DATAPSP - EBOOT_SND0 > 0:
			EBOOT_EXTRACT = open(EBOOT_DEST + os.sep + "SND0.AT3","wb")
			EBOOT_FILE.seek(EBOOT_SND0)
			FILE_SIZE = EBOOT_DATAPSP - EBOOT_SND0
			while FILE_SIZE > 0:
				if FILE_SIZE >= 1024 * 256:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(1024 * 256))
					FILE_SIZE = FILE_SIZE - 1024 * 256
				else:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(FILE_SIZE))
					FILE_SIZE = 0
			EBOOT_EXTRACT.close()
		if EBOOT_DATAPSAR - EBOOT_DATAPSP > 0:
			EBOOT_EXTRACT = open(EBOOT_DEST + os.sep + "DATA.PSP","wb")
			EBOOT_FILE.seek(EBOOT_DATAPSP)
			FILE_SIZE = EBOOT_DATAPSAR - EBOOT_DATAPSP
			while FILE_SIZE > 0:
				if FILE_SIZE >= 1024 * 256:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(1024 * 256))
					FILE_SIZE = FILE_SIZE - 1024 * 256
				else:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(FILE_SIZE))
					FILE_SIZE = 0
			EBOOT_EXTRACT.close()
		if EBOOT_END - EBOOT_DATAPSAR > 0:
			EBOOT_EXTRACT = open(EBOOT_DEST + os.sep + "DATA.PSAR","wb")
			EBOOT_FILE.seek(EBOOT_DATAPSAR)
			FILE_SIZE = EBOOT_END - EBOOT_DATAPSAR
			while FILE_SIZE > 0:
				if FILE_SIZE >= 1024 * 256:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(1024 * 256))
					FILE_SIZE = FILE_SIZE - 1024 * 256
				else:
					EBOOT_EXTRACT.write(EBOOT_FILE.read(FILE_SIZE))
					FILE_SIZE = 0
			EBOOT_EXTRACT.close()
		return True
	else:
		return False
#End of unpack_eboot
